import backtrader as bt

class SuperTrend(bt.Indicator):
    params = (
        ('period', 10),      # ATR计算周期（默认10）
        ('multiplier', 3),   # ATR乘数（默认3）
    )
    lines = ('supertrend', 'trend', 'supertrend_up', 'supertrend_down') 
    plotinfo = dict(subplot=False)   # 与主图同画布
    # 定义 plotlines，动态设置颜色
    plotlines = dict(
        supertrend=dict(
            color='green',   # 默认颜色（会被覆盖）
            ls='-',          # 实线
            _plotskip=True,  # 禁用默认绘图（手动控制）
        ),
        # 新增两条线，分别用于上涨和下跌
        supertrend_up=dict(
            color='green',
            ls='-',
        ),
        supertrend_down=dict(
            color='red',
            ls='-',
        ),
    )

    def __init__(self):
        # 计算中间价：(high + low) / 2
        self.mid = (self.data.high + self.data.low) / 2
        # 计算ATR
        self.atr = bt.indicators.ATR(self.data, period=self.p.period)
        # 计算上下通道
        self.upper = self.mid + self.p.multiplier * self.atr
        self.lower = self.mid - self.p.multiplier * self.atr

        # 初始化趋势方向（默认空头）
        self.trend_dir = -1

    def next(self):
        # 第一根K线初始化
        if len(self) <= self.p.period+1:
            self.lines.supertrend[0] = self.lower[0]
            self.lines.trend[0] = self.trend_dir
            return

        # 获取前一根的趋势方向和超级趋势值
        prev_trend = self.lines.trend[-1]
        prev_supertrend = self.lines.supertrend[-1]
        # 趋势判断逻辑
        if prev_trend == 1:  # 当前多头趋势
            if self.data.close[0] < prev_supertrend:  # 价格跌破超级趋势线
                self.trend_dir = -1
                self.lines.supertrend[0] = self.upper[0]  # 切换为下通道
            else:
                self.trend_dir = 1
                self.lines.supertrend[0] = max(prev_supertrend, self.lower[0])  # 防止跳空

        else:  # 当前空头趋势
            if self.data.close[0] > prev_supertrend:  # 价格突破超级趋势线
                self.trend_dir = 1
                self.lines.supertrend[0] = self.lower[0]  # 切换为上通道
            else:
                self.trend_dir = -1
                self.lines.supertrend[0] = min(prev_supertrend, self.upper[0])  # 防止跳空


        self.lines.trend[0] = self.trend_dir
        # 动态更新两条线
        if self.trend_dir == 1:
            self.lines.supertrend_up[0] = self.lines.supertrend[0]
            self.lines.supertrend_down[0] = float('nan')  # 隐藏下跌线
        else:
            self.lines.supertrend_down[0] = self.lines.supertrend[0]
            self.lines.supertrend_up[0] = float('nan')   # 隐藏上涨线